import os
import tensorflow as tf

class Cifar(object):
    def __init__(self):
        # 初始化操作
        self.height = 32
        self.width = 32
        self.channels = 3

        # 字节数
        self.image_bytes = self.height * self.width * self.channels
        self.label_bytes = 1
        self.all_bytes = self.label_bytes + self.image_bytes

    def read_and_decode(self,file_list):
        """
        二进制文件读取。将图片读取成image_batch,将label读取成label_batch
        :param file_list:
        :return:
        """
        # 1 构造文件队列
        file_queue = tf.train.string_input_producer(file_list)
        # 2 读取与解码
        reader = tf.FixedLengthRecordReader(self.all_bytes)
        # key文件名， value一个数据样本。
        key,value = reader.read((file_name))
        # 解码
        decoded = tf.decode_raw(value,tf.uint8)
        # 将label与data分开第一个字节是label,后面3072个是图片
        label = tf.slice(decoded,[0],[self.label_bytes])
        image = tf.slice(decoded,[self.label_bytes],[self.image_bytes])

        # 调整图片大小 动态形状
        image_reshaped = tf.reshape(image,shape=[self.channels,self.height,self.width])

        # 转置 将图片的顺序转为height,width,channels
        image_transposed = tf.transpose(image_reshaped,[1,2,0])

        # 调整图像类型 一般存储用uint8 , 计算时用float32
        image_cast = tf.cast(image_transposed,tf.float32)
        # 3 批处理
        label_batch,image_batch = tf.train.batch([label,image_cast],batch_size=100,num_threads=1,capacity=100)
        # 开启会话
        with tf.Session() as sess:
            # 开启线程
            coord = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(sess,coord)

            # key_new,value_new,decoded,label_batch_new,image_batch = sess.run([key,value,decoded,label_batch,image_batch])
            # print("key_new: ", key_new)
            # print("value_new: ", value_new)
            # print("decoded: ", decoded)
            # print("label_batch_new: ", label_batch_new)
            # print("image_batch: ", image_batch)
            label_value,image_value = sess.run([label_batch,image_batch])

            #回收线程
            coord.request_stop()
            coord.join(threads)

        return image_value,label_value

    def write_to_tfrecords(self,image_batch,label_batch):
        """
        将样本的特征值和目标值一起写入tfrecords文件。
        :param image:
        :param label:
        :return:
        """

        with tf.python_io.TFRecordWriter("cifar10.tfrecords") as writer:
            #循环构造example对象，并序列化写入文件：
            for i in range(100):
                image = image_batch[i].tostring()
                label = label_batch[i][0]
            example = tf.train.Example(features=tf.train.Features(feature={
                "image":tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])),
                "label":tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
            }))
            writer.write(example.SerializeToString())
        return None

    def read_tfrecords(self):
        """
            读取tfrecords文件
        :return:
        """
        # 1 构建文件名列表
        file_queue = tf.train.string_input_producer(["cifar10.tfrecords"])
        # 2 读取与解码
        # 读取
        reader = tf.TFRecordReader()
        # key文件名， value 读到的一个example样本
        key,value = reader.read(file_queue)
        # 解析example
        feature = tf.parse_single_example(value,features={
            "image":tf.FixedLenFeature([],tf.string),
            "label":tf.FixedLenFeature([],tf.int64)
        })
        image = feature["image"]
        label = feature["label"]
        # 解码
        image_decoded = tf.decode_raw(image,out_type=tf.uint8)
        # 图像形状调整
        image_reshaped = tf.reshape(image_decoded,shape=[self.height,self.width,self.channels])
        print("image_reshaped : ",image_reshaped)

        # 3 构造批处理队列
        image_batch,label_batch = tf.train.batch([image_reshaped,label],batch_size=100,num_threads=2,capacity=100)
        # 开启会话
        with tf.Session() as sess:
            coord = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(sess,coord=coord)

            image_value , label_value= sess.run([image_batch,label_batch])
            print("image_value : ",image_value)
            print("label_value : ", label_value)

            coord.request_stop()
            coord.join(threads)
        return None

if __name__ == "__main__":
    file_name = os.listdir("cifar-10-batches-bin")
    #构造文件名路径列表
    file_list = [os.path.join("./cifar-10-batches-bin/",file) for file in file_name if file[-3:]=="bin"]

    #实例化对象
    cifar = Cifar()
    image_batch, label_batch = cifar.read_and_decode(file_list)
    cifar.write_to_tfrecords(image_batch,label_batch)
